
import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
from pev1.StockPool import  stockPool
import traceback

"""
判断某只股票在某日是否满足K线上穿10日均线
:param code: 股票代码
:param beginDate: 开始日期
:param endDate: 结束日期
:return: True/False
"""
def isKUpBreakMA10(code, beginDate, endDate):
    # 如果没有指定开始日期和结束日期，则直接返回False
    if beginDate is None or endDate is None:
        return False

    # 如果股票当日停牌或者是下跌，则返回False
    dailyKS = ts.get_k_data(code,autype='hfq', start=beginDate, end=endDate)

    # 需要判断两日的K线和10日均线的相对位置，所以如果K线数不满足11个，也就是无法计算两个MA10,周日直接返回False
    indexSize = dailyKS.index.size
    if indexSize < 11:
        return False

    # 计算MA10
    dailyKS['ma'] = dailyKS['close'].rolling(10).mean()

    # 比较收盘价和MA10的关系
    dailyKS['delta'] = dailyKS['close']  - dailyKS['ma']

    return  dailyKS.loc[dailyKS.index[9]]['delta'] <= 0 < dailyKS.loc[dailyKS.index[10]]['delta']

"""
判断某只股票在某日是否满足K线下穿10日均线
:param code: 股票代码
:param beginDate: 开始日期
:param endDate: 结束日期
:return: True/False
"""
def isKDownBreakMA10(code, beginDate, endDate):

    # 如果没有指定开始日期和结束日期，则直接返回False
    if beginDate is None or endDate is None:
        return False

    dailyKS = ts.get_k_data(code, start=beginDate, end=endDate)

    # 需要判断两日的K线和10日均线的相对位置，所以如果K线数不满足11个，也就是无法计算两个MA10,则直接返回False
    if dailyKS.index.size < 11:
        return False

    # 计算MA10
    dailyKS['ma'] = dailyKS['close'].rolling(10).mean()

    # 计算收盘价与MA10的差
    dailyKS['delta'] = dailyKS['close'] - dailyKS['ma']

    return dailyKS.loc[dailyKS.index[9]]['delta'] >= 0 > dailyKS.loc[dailyKS.index[10]]['delta']

"""
计算年化收益
"""
def computeAnnualProfit(tradingDays, netValue):
    annualProfit = 0
    if tradingDays > 0:
        # 计算年数
        years = tradingDays / 245

        # 计算年化收益
        annualProfit = pow(netValue, 1 / years) - 1

    annualProfit = round(annualProfit * 100, 2)

    return annualProfit



"""
计算夏普比率
:param netValues: 净值列表
"""
def computeSharpeRatio(netValues):
    # 总交易日数
    tradingDays = len(netValues)

    # 所有收益的DataFrame
    profitDF = pd.DataFrame(columns={'profit'})

    # 收益之后，初始化为第一天的收益
    profitDF.loc[0] = {'profit': round((netValues[0] - 1) * 100, 2)}

    # 计算每天的收益
    for index in range(1, tradingDays):
        # 计算每日的收益变化
        profit = (netValues[index] - netValues[index - 1]) / netValues[index - 1]
        profit = round(profit * 100, 2)
        profitDF.loc[index] = {'profit': profit}

    # 计算标准差
    profitSTD = pow(profitDF.var()['profit'], 1 / 2)

    # 计算年化收益
    annualProfit = computeAnnualProfit(tradingDays, netValues[-1])

    # 夏普比率
    sharpeRatio = (annualProfit - 4.75) / profitSTD

"""
计算最大回撤
:param netValues: 净值列表
"""
def computeDrawDown(netValues):
    # 最大回护初始值设为0
    maxDrawDown = 0

    size = len(netValues)

    # 双层循环找出最大回撤
    index = 0
    for netValue in netValues:
        for subNet in netValues[index:]:
            drawDown = 1 - subNet / netValue
            if drawDown > maxDrawDown:
                maxDrawDown = drawDown
        index += 1

    return  maxDrawDown


"""
找到上期入选本期被调出的股票，这些股票将必须卖出
:param lastPhaseCodes: 上期的股票列表
:param thisPhaseCodes: 本期的股票列表
:return: 被调出的股票列表
"""
def findOutStocks(lastPhaseCodes, thisPhaseCodes):
    outStocks = []

    for code in lastPhaseCodes:
        if code not in thisPhaseCodes:
            outStocks.append(code)

    return outStocks


"""
策略回测。结束后打印出收益曲线(沪深300基准)、年化收益、最大回撤、

:param beginDate: 回测开始日期
:param endDate: 回测结束日期
"""
def backTest(beginDate, endDate):
    cash = 1E7
    singlePosition = 2E5

    # 保存持仓股的日期
    codeDateVolumeDict = dict()

    # 时间为key的净值、收益和同期沪深基准
    dfProfit = pd.DataFrame(columns={'net_value', 'profit', 'hs300'})

    # 因为上证指数没有停牌不会缺数， 所以用它作为交易日历
    szzsDF = ts.get_k_data('000001', index=True, start=beginDate, end=endDate)
    allDates = list(szzsDF['date'])

    # 获取沪深300在统计周期内的第一天的值
    hs300K = ts.get_k_data('000300', index=True, start=allDates[0], end=allDates[0])
    hs300BeginValue = hs300K.loc[hs300K.index[0]]['close']

    # 获取股票数据
    rebalanceDates, dateCodesDict = stockPool(beginDate, endDate)

    # 获取回测周期内股票池内所有股票的收盘价和前收价
    allOptionCodeSet = set()

    for rebalanceDate in rebalanceDates:
        for code in dateCodesDict[rebalanceDate]:
            allOptionCodeSet.add(code)

    # 缓存股票的日线数据
    codeDailyDict = dict()
    for code in allOptionCodeSet:
        dailiesDF = ts.get_k_data(code, autype=None, start=beginDate, end=endDate)
        dailiesHfqDF = ts.get_k_data(code, autype='hfq', start=beginDate, end=endDate)

        # 计算复权因子
        dailiesDF['au_factor'] = dailiesHfqDF['close'] / dailiesDF['close']
        dailiesDF.set_index(['date'], inplace=True)

        codeDailyDict[code] = dailiesDF

    lastPhaseCodes = None
    thisPhaseCodes = None
    toBeSoldCodes = set()
    toBeBoughtCodes = set()
    holdingCodeDict = dict()
    lastDate = None

    # 按照日期一步步回测
    for date in allDates:
        print('BackTest at %s.' % date )

        # 当期持仓股票列表
        # {'600000': {'cost': 500, 'volume': 000, 'last_value': 234}}
        beforeSellHoldingCodes = list(holdingCodeDict.keys())

        # 处理持仓股除权除息
        if lastDate is not None and len(beforeSellHoldingCodes) > 0:
            for code in beforeSellHoldingCodes:
                try:
                    dailies = codeDailyDict[code]
                    # 上一个交易日的复权因子
                    currentAuFactor = dailies.loc[date]['au_factor']
                    beforeVolume = holdingCodeDict[code]['volume']
                    lastAuFactor = dailies[lastDate]['au_factor']

                    afterVolume = int(beforeVolume * (currentAuFactor / lastAuFactor))
                    holdingCodeDict[code]['volume'] = afterVolume
                    print('持仓量调整：%s, %6d, %10.6f, %6d, %10.6f' %
                      (code, beforeVolume, lastAuFactor, afterVolume, currentAuFactor), flush=True)
                except:
                    print('持仓量调整时，发生错误：%s, %s' % (code, date), flush=True)
                    traceback.print_exc()

        # 买入
        if len(toBeBoughtCodes) > 0:
            for code in toBeBoughtCodes:
                try:
                    if cash > singlePosition:
                        buyPrice = codeDailyDict[code].loc[date]['open']
                        volume = int(int(singlePosition / buyPrice) / 100) * 100
                        buyAmount = buyPrice * volume
                        cash -= buyAmount
                        holdingCodeDict[code] = {
                            'volume': volume,
                            'cost': buyAmount,
                            'last_value': buyAmount
                        }
                        print('买入 %s, %6d, %6.2f, %8.2f' % (code, volume, buyPrice, buyAmount), flush=True)
                except:
                    print('买入时，发生错误：%s, %s' % (code, date), flush=True)
                    traceback.print_exc()
        print('买入后，现金: %10.2f' % cash)

        # 卖出
        if len(toBeSoldCodes) > 0:
            codeSetTmp = set(toBeSoldCodes)
            for code in codeSetTmp:
                try:
                    if code in beforeSellHoldingCodes:
                        holdingStock = holdingCodeDict[code]
                        holdingVolume = holdingStock['volume']
                        sellPrice = codeDailyDict[code].loc[date]['open']
                        sellAmount = sellPrice * holdingVolume
                        cash += sellAmount

                        cost = holdingStock['cost']
                        singleProfit = (sellAmount - cost) * 100 / cost
                        print('卖出 %s, %6d, %6.2f, %8.2f, %4.2f' %
                              (code, holdingVolume, sellPrice, sellAmount, singleProfit))

                        del holdingCodeDict[code]
                        toBeSoldCodes.remove(code)
                except:
                    print('卖出时，发生异常：%s, %s' % (code, date), flush=True)
                    traceback.print_exc()
        print('卖出后，现金: %10.2f' % cash)

        # 持仓股代码列表
        holdingCodes = list(holdingCodeDict.keys())

        # 如果调整日， 则获取新一期的股票列表
        if date in rebalanceDates:
            # 暂存为上期的日期
            if thisPhaseCodes is not None:
                lastPhaseCodes = thisPhaseCodes
            thisPhaseCodes = dateCodesDict[date]

            # 找到所有调出股票代码，在第二日开盘时卖出
            if lastPhaseCodes is not None:
                outCodes = findOutStocks(lastPhaseCodes, thisPhaseCodes)
                for outCode in outCodes:
                    if outCode in holdingCodeDict:
                        toBeSoldCodes.add(outCode)

        # 获取检测信号的开始日期和结束日期
        currentDateIndex = allDates.index(date)
        signalBeginDate = None
        if currentDateIndex >= 11:
            signalBeginDate = allDates[currentDateIndex - 11]

        # 检查是否有需要第二天卖出的股票
        for holdingCode in holdingCodes:
            if isKDownBreakMA10(holdingCode, beginDate= signalBeginDate, endDate= date):
                toBeSoldCodes.add(holdingCode)

        # 检查是否需要第二天买入的股票
        toBeBoughtCodes.clear()
        if thisPhaseCodes is not None:
            for code in thisPhaseCodes:
                if code not in holdingCodes and isKDownBreakMA10(code, beginDate=signalBeginDate, endDate=date):
                    toBeBoughtCodes.add(code)
        # 计算总资产
        totalValue = 0
        for code in holdingCodes:
            try:
                holdingStock = holdingCodeDict[code]
                value = codeDailyDict[code].loc[date]['close'] * holdingStock['volume']
                totalValue += value

                # 计算总收益
                profit = (value - holdingStock['cost']) * 100 / holdingStock['cost']

                # 计算当日收益
                oneDayProfit = (value - holdingStock['last_value']) * 100 / holdingStock['last_value']

                # 暂存当日市值
                holdingStock['last_value'] = value

                print('持仓: %s, %10.2f, %4.2f, %4.2f' %
                      (code, value, profit, oneDayProfit))

                # 保存每一日股票的持仓数
                codeDateVolumeDict[code + '_' + date] = holdingStock['volume']
            except:
                print('计算收益时发生错误：%s, %s' % (code, date), flush=True)
                traceback.print_exc()
        totalCapital = totalValue + cash

        # 计算基准收益
        hs300KCurrent = ts.get_k_data('000300', index=True, start=date, endDate=date)
        hs300CurrentValue = hs300KCurrent.loc[hs300KCurrent.index[0]]['close']

        print('收盘后，现金: %10.2f, 总资产: %10.2f' % (cash, totalCapital), flush=True)
        lastDate = date
        dfProfit.loc[date] = {
            'net_value': round(totalCapital / 1e7, 2),
            'profit': round(100 * (totalCapital - 1e7) /1e7, 2),
            'hs300': round(100 * (hs300CurrentValue - hs300BeginValue) / hs300BeginValue, 2)
        }

    # 打印回测收益曲线数值
    for indexDate in dfProfit.index:
        print('%s, %6.2f, %6.2f' %
              (indexDate, dfProfit.loc[indexDate]['profit'], dfProfit.loc[indexDate]['hs300']),
              flush=True)
    drawDown = computeDrawDown(dfProfit['net_value'])
    annualProfit, sharpeRatio = computeSharpeRatio(dfProfit['net_value'])

    print('回测结果 %s - %s，年化收益： %7.3f, 最大回撤：%7.3f, 夏普比率：%4.2f' %
          (beginDate, endDate, annualProfit, drawDown, sharpeRatio))

    dfProfit.plot(title='BackTest Result', y = ['profit', 'hs300'], kind = 'line')
    plt.show()

if __name__ == '__main__':

    df = pd.DataFrame([10,11,12],columns=['profit'])
    print(df.var()['profit'])



